import numpy as np
import pandas as pd
from sklearn.tree import DecisionTreeRegressor, plot_tree
import matplotlib.pyplot as plt
import random
from sklearn.metrics import mean_absolute_error, r2_score
from sklearn.model_selection import GridSearchCV, cross_val_score
from sklearn.preprocessing import StandardScaler
import copy
import os

# 读取Excel文件中的前70个子表，再读取后5个表作为新测试样本
dfs = pd.read_excel('Bottleneck_100.xlsx', sheet_name=None, engine='openpyxl')
table_names = list(dfs.keys())[25:95]
table_new_names = list(dfs.keys())[95:100]

# 随机划分80%作为训练集，20%作为测试集，新测试样本不动
random.shuffle(table_names)
train_size = int(len(table_names) * 0.8)
train_tables = table_names[:train_size]
test_tables = table_names[train_size:]
new_test_tables = table_new_names

# 训练基模型
models = []
scalers = {}
for table in train_tables:
    df = dfs[table]
    X = df.iloc[:, :8]  # 前8列是特征
    y = df.iloc[:, 8]  # 第9列是标签

    # 特征归一化
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    scalers[table] = scaler

    model = DecisionTreeRegressor(max_depth=5, min_samples_split=3, min_samples_leaf=2)
    model.fit(X, y)
    models.append(model)

    # 指定保存图片的文件夹路径
    output_folder_tree = "tree_picture"
    if not os.path.exists(output_folder_tree):
        os.makedirs(output_folder_tree)
    # 对模型做深拷贝，仅用于绘图，避免影响原模型预测
    viz_model = copy.deepcopy(model)
    tree = viz_model.tree_
    # 对每个非叶子节点的阈值进行反归一化
    for i in range(tree.node_count):
        if tree.feature[i] != -2:  # -2 表示叶子节点
            feature_index = tree.feature[i]
            # 根据StandardScaler的公式进行反变换
            tree.threshold[i] = tree.threshold[i] * scaler.scale_[feature_index] + scaler.mean_[feature_index]
    # 使用反变换后的模型绘制决策树
    plt.figure(figsize=(15, 8))
    plot_tree(viz_model, feature_names=df.columns[:8], filled=True)
    plt.title(f'Decision Tree for {table}')
    # 构造保存图片的完整路径（这里用表名作为文件名的一部分）
    file_path_tree = os.path.join(output_folder_tree, f"DecisionTree_{table}.png")
    plt.savefig(file_path_tree)  # 保存图片到指定文件夹
    plt.close()  # 关闭当前图形，防止重复绘制


# 预测函数（基于多个基模型的平均预测）
def ensemble_predict(models, scalers, X):
    predictions = np.array([model.predict(scalers[table].transform(X))
                            for model, table in zip(models, train_tables)])
    return np.mean(predictions, axis=0)

# 计算测试集的预测误差，并绘制实际值与预测值对比图
error_metrics = []
for table in test_tables:
    df = dfs[table]
    X_test = df.iloc[:, :8]
    y_test = df.iloc[:, 8]
    sample_names = df.iloc[:, 9]  # 第10列是工序编号

    y_pred = ensemble_predict(models, scalers, X_test)
    mae = mean_absolute_error(y_test, y_pred)
    r2 = r2_score(y_test, y_pred)
    error_metrics.append((rmse, mae, r2))

    print(f'表 {table} - MAE: {mae:.4f}, R2: {r2:.4f}')

    # 绘制实际值与预测值对比图
    plt.figure(figsize=(12, 6))
    plt.rcParams['font.family'] = 'SimHei'
    plt.plot(sample_names, y_test, label='实际瓶颈值', marker='o')
    plt.plot(sample_names, y_pred, label='预测瓶颈值', marker='s')
    plt.xlabel('工序编号')
    plt.ylabel('瓶颈值')
    plt.title(f'第{table}次生产活动下的实际瓶颈值与预测瓶颈值对比')
    plt.legend()
    plt.xticks(rotation=45)
    plt.grid()

    output_folder_test = "test_picture"
    if not os.path.exists(output_folder_test):
        os.makedirs(output_folder_test)
    file_path_test = os.path.join(output_folder_test, f"test_loss_{table}.png")
    plt.savefig(file_path_test)  # 保存图片到指定文件夹
    plt.close()  # 关闭当前图形，防止重复绘制

# 计算整体误差
average_mae = np.mean([m[1] for m in error_metrics])
average_r2 = np.mean([m[2] for m in error_metrics])
print(f'测试集平均 MAE: {average_mae:.4f}, R2: {average_r2:.4f}')

print('--------')

# 计算新测试集的预测误差，并绘制实际值与预测值对比图
error_metrics = []
for table in new_test_tables:
    df = dfs[table]
    X_test = df.iloc[:, :8]
    y_test = df.iloc[:, 8]
    sample_names = df.iloc[:, 9]  # 第10列是工序编号

    y_pred = ensemble_predict(models, scalers, X_test)
    mae = mean_absolute_error(y_test, y_pred)
    r2 = r2_score(y_test, y_pred)
    error_metrics.append((rmse, mae, r2))

    print(f'表 {table} - RMSE: {rmse:.4f}, MAE: {mae:.4f}, R2: {r2:.4f}')

    # 绘制实际值与预测值对比图
    plt.figure(figsize=(12, 6))
    plt.rcParams['font.family'] = 'SimHei'
    plt.plot(sample_names, y_test, label='实际瓶颈值', marker='o')
    plt.plot(sample_names, y_pred, label='预测瓶颈值', marker='s')
    plt.xlabel('工序编号')
    plt.ylabel('瓶颈值')
    plt.title(f'第{table}次生产活动下的实际瓶颈值与预测瓶颈值对比')
    plt.legend()
    plt.xticks(rotation=45)
    plt.grid()

    output_folder_new_test = "new_test_picture"
    if not os.path.exists(output_folder_new_test):
        os.makedirs(output_folder_new_test)
    file_path_new_test = os.path.join(output_folder_new_test, f"new_test_loss_{table}.png")
    plt.savefig(file_path_new_test)  # 保存图片到指定文件夹
    plt.close()  # 关闭当前图形，防止重复绘制

# 计算整体误差
average_mae = np.mean([m[1] for m in error_metrics])
average_r2 = np.mean([m[2] for m in error_metrics])
print(f'新测试集平均MAE: {average_mae:.4f}, R2: {average_r2:.4f}')